// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/functions/lib/repeat.h>
#include <pollux/testing/gtest_utils.h>
#include <pollux/common/testutil/optional_empty.h>
#include <pollux/testing/functions/prestosql/function_base_test.h>

using namespace kumo::pollux::test;

namespace kumo::pollux::functions {
namespace {

class RepeatTest : public functions::test::FunctionBaseTest {
 protected:
  static void SetUpTestCase() {
    FunctionBaseTest::SetUpTestCase();
    exec::registerStatefulVectorFunction(
        "repeat",
        functions::repeatSignatures(),
        functions::makeRepeat,
        functions::repeatMetadata());
    exec::registerStatefulVectorFunction(
        "repeat_allow_negative_count",
        functions::repeatSignatures(),
        functions::makeRepeatAllowNegativeCount,
        functions::repeatMetadata());
  }

  void testExpression(
      const std::string& expression,
      const std::vector<VectorPtr>& input,
      const VectorPtr& expected) {
    auto result = evaluate(expression, make_row_vector(input));
    assertEqualVectors(expected, result);
  }

  void testExpressionWithError(
      const std::string& expression,
      const std::vector<VectorPtr>& input,
      const std::string& expectedError) {
    POLLUX_ASSERT_THROW(
        evaluate(expression, make_row_vector(input)), expectedError);
  }
};

TEST_F(RepeatTest, repeat) {
  const auto elementVector = make_nullable_flat_vector<float>(
      {0.0, -2.0, 3.333333, 4.0004, std::nullopt, 5.12345});
  const auto countVector =
      make_nullable_flat_vector<int32_t>({1, 2, 3, 0, 4, std::nullopt});
  VectorPtr expected;

  expected = make_nullable_array_vector<float>({
      {{0.0}},
      {{-2.0, -2.0}},
      {{3.333333, 3.333333, 3.333333}},
      common::testutil::optionalEmpty,
      {{std::nullopt, std::nullopt, std::nullopt, std::nullopt}},
      std::nullopt,
  });
  testExpression("repeat(C0, C1)", {elementVector, countVector}, expected);
  testExpression("try(repeat(C0, C1))", {elementVector, countVector}, expected);

  // Test using a null constant as the count argument.
  expected = BaseVector::create_null_constant(ARRAY(REAL()), 6, pool());
  testExpression("repeat(C0, null::INTEGER)", {elementVector}, expected);
  testExpression("try(repeat(C0, null::INTEGER))", {elementVector}, expected);

  // Test using a non-null constant as the count argument.
  expected = make_nullable_array_vector<float>({
      {0.0, 0.0, 0.0},
      {-2.0, -2.0, -2.0},
      {3.333333, 3.333333, 3.333333},
      {4.0004, 4.0004, 4.0004},
      {std::nullopt, std::nullopt, std::nullopt},
      {5.12345, 5.12345, 5.12345},
  });
  testExpression("repeat(C0, '3'::INTEGER)", {elementVector}, expected);
  testExpression("try(repeat(C0, '3'::INTEGER))", {elementVector}, expected);

  expected = make_array_vector<float>({{}, {}, {}, {}, {}, {}});
  testExpression("repeat(C0, '0'::INTEGER)", {elementVector}, expected);
  testExpression("try(repeat(C0, '0'::INTEGER))", {elementVector}, expected);
}

TEST_F(RepeatTest, repeatWithInvalidCount) {
  const auto elementVector =
      make_nullable_flat_vector<float>({0.0, 2.0, 3.333333});

  VectorPtr countVector;
  VectorPtr expected;

  expected = make_nullable_array_vector<float>({
      {{0.0}},
      std::nullopt,
      {{3.333333, 3.333333, 3.333333}},
  });
  countVector = make_nullable_flat_vector<int32_t>({1, -2, 3});
  testExpression("try(repeat(C0, C1))", {elementVector, countVector}, expected);
  testExpressionWithError(
      "repeat(C0, C1)",
      {elementVector, countVector},
      "(-2 vs. 0) Count argument of repeat function must be greater than or equal to 0");

  countVector = make_nullable_flat_vector<int32_t>({1, 123456, 3});
  testExpression("try(repeat(C0, C1))", {elementVector, countVector}, expected);
  testExpressionWithError(
      "repeat(C0, C1)",
      {elementVector, countVector},
      "(123456 vs. 10000) Count argument of repeat function must be less than or equal to 10000");

  // Test using a constant as the count argument.
  expected = BaseVector::create_null_constant(ARRAY(REAL()), 3, pool());
  testExpression("try(repeat(C0, '-5'::INTEGER))", {elementVector}, expected);
  testExpressionWithError(
      "repeat(C0, '-5'::INTEGER)",
      {elementVector},
      "(-5 vs. 0) Count argument of repeat function must be greater than or equal to 0");

  testExpression(
      "try(repeat(C0, '10001'::INTEGER))", {elementVector}, expected);
  testExpressionWithError(
      "repeat(C0, '10001'::INTEGER)",
      {elementVector},
      "(10001 vs. 10000) Count argument of repeat function must be less than or equal to 10000");
}

TEST_F(RepeatTest, repeatAllowNegativeCount) {
  const auto elementVector = make_nullable_flat_vector<float>(
      {0.0, -2.0, 3.333333, 4.0004, std::nullopt, 5.12345});
  auto expected = make_array_vector<float>({{}, {}, {}, {}, {}, {}});

  // Test negative count.
  auto countVector =
      make_nullable_flat_vector<int32_t>({-1, -2, -3, -5, -10, -100});
  testExpression(
      "repeat_allow_negative_count(C0, C1)",
      {elementVector, countVector},
      expected);

  // Test using a constant as the count argument.
  testExpression(
      "repeat_allow_negative_count(C0, '-5'::INTEGER)",
      {elementVector},
      expected);

  // Test mixed case.
  expected = make_array_vector<float>(
      {{0.0}, {-2.0, -2.0}, {}, {}, {}, {5.12345, 5.12345, 5.12345}});
  countVector = make_nullable_flat_vector<int32_t>({1, 2, -1, 0, -10, 3});
  testExpression(
      "repeat_allow_negative_count(C0, C1)",
      {elementVector, countVector},
      expected);
}
} // namespace
} // namespace kumo::pollux::functions
